# univariate convlstm example
from keras.layers import ConvLSTM2D
from numpy import array
import matplotlib.pyplot as plt
from pandas import read_csv
from keras.models import Sequential
from keras.layers import LSTM
from keras.layers import Dense
from keras.layers import Flatten
from math import sqrt
from sklearn.preprocessing import MinMaxScaler
from keras.callbacks import EarlyStopping
import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score


def create_dataset(dataset, scale):
    train_size = int(len(dataset) * scale)
    # 将数据集分为两部分
    train, test = dataset[0:train_size], dataset[train_size:]
    # train, test = dataset.iloc[0:train_size], dataset.iloc[train_size:]

    return train, test


# split a univariate sequence into samples
def split_sequence(sequence, n_steps):
	X, y = list(), list()
	for i in range(len(sequence)):
		# find the end of this pattern
		end_ix = i + n_steps
		# check if we are beyond the sequence
		if end_ix > len(sequence)-1:
			break
		# gather input and output parts of the pattern
		seq_x, seq_y = sequence[i:end_ix], sequence[end_ix]
		X.append(seq_x)
		y.append(seq_y)
	return array(X), array(y)


def mean_absolute_percentage_error(y_true, y_pred):
	# 平均绝对百分比误差（MAPE）的计算
	y_true, y_pred = np.array(y_true), np.array(y_pred)
	return np.mean(np.abs((y_true - y_pred) / y_true)) * 100


# evaluate the RMSE for each forecast time step
def evaluate_forecasts(data, forecasts):
	rmse = sqrt(mean_squared_error(data, forecasts))
	mae = mean_absolute_error(data, forecasts)
	mape = mean_absolute_percentage_error(data, forecasts)
	r2 = r2_score(data, forecasts)
	print(' RMSE: %f' % rmse)
	print(' MAE: %f' % mae)
	print(' MAPE: %f' % mape)
	print(' R2: %f' % r2)

	return rmse,mae,mape,r2

if __name__ == '__main__':
	# 加载数据
	dataframe = read_csv('E:\lyf_ML_Drought\coding\ML_Drought_Prediction\indices_caculate\\result\ROW_SPEI\ROW_SPEI-12\SPEI-12_52533.txt', header=None, names=('TIME','SPEI-1'))
	dataframe = dataframe.set_index(['TIME'], drop=True) # 把日期作为索引
	# print(dataframe.values.T.tolist()[0])
	# dataset = dataframe.values.T.tolist()[0]

	# 创建训练和测试数据集
	train, test = create_dataset(np.array(dataframe), 0.9)

	# choose a number of time steps
	time_steps = 4
	# split into samples
	X, y = split_sequence(train, time_steps)
	# reshape from [samples, timesteps] into [samples, subsequences, timesteps, features]
	n_features = 1
	n_seq = 2
	n_steps = 2
	X = X.reshape((X.shape[0], n_seq, 1, n_steps, n_features))
	# define model
	model = Sequential()
	model.add(
		ConvLSTM2D(filters=64, kernel_size=(1, 2), activation='relu', input_shape=(n_seq, 1, n_steps, n_features)))
	model.add(Flatten())
	model.add(Dense(1))
	model.compile(optimizer='adam', loss='mean_squared_error')
	# fit model
	model.fit(X, y, epochs=3000,batch_size=50, verbose=1, shuffle=False)
	x_input, y_input = split_sequence(test, time_steps)
	x_input = x_input.reshape((x_input.shape[0], n_seq, 1, n_steps, n_features))

	train_pred = model.predict(X, verbose=0)
	yhat = model.predict(x_input, verbose=0)
	# print(yhat)

	print("训练模型的评估：")
	evaluate_forecasts(y, train_pred)
	print("模型测试的评估：")
	evaluate_forecasts(y_input, yhat)

	plt.figure()
	plt.plot(y_input, label='True', c="black")
	plt.plot(yhat, label='LSTM', c="red")
	plt.title("ConvLSTM's Prediction")
	plt.xlabel('Month')
	plt.ylabel('SPEI')
	plt.legend()
	plt.show();